# -*- coding: utf-8 -*-
"""
主脚本

整合所有功能，提供完整的工作流程，包括数据处理、模型训练、遗传算法优化和SHAP分析。
"""

import os
import sys
import argparse
import time
import json
import matplotlib.pyplot as plt

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


def create_directories():
    """
    创建必要的目录
    """
    directories = ['output', 'models', 'figures', 'ga_results', 'shap_results']
    for directory in directories:
        if not os.path.exists(directory):
            os.makedirs(directory)
            print(f"创建目录: {directory}")


def run_data_processing():
    """
    运行数据处理脚本
    """
    print("\n" + "=" * 50)
    print("开始数据处理...")
    print("=" * 50)
    
    import data_processing
    data_processing.main()


def run_model_training():
    """
    运行模型训练脚本
    """
    print("\n" + "=" * 50)
    print("开始模型训练...")
    print("=" * 50)
    
    import train
    train.main()


def run_genetic_algorithm():
    """
    运行遗传算法脚本
    """
    print("\n" + "=" * 50)
    print("开始遗传算法优化...")
    print("=" * 50)
    
    import genetic_algorithm
    genetic_algorithm.main()


def run_shap_analysis():
    """
    运行SHAP分析脚本
    """
    print("\n" + "=" * 50)
    print("开始SHAP分析...")
    print("=" * 50)
    
    import shap_analysis
    shap_analysis.main()


def generate_summary_report():
    """
    生成总结报告
    """
    print("\n" + "=" * 50)
    print("生成总结报告...")
    print("=" * 50)
    
    # 创建报告内容
    report = []
    
    # 添加标题
    report.append("# 材料性能预测系统总结报告")
    report.append("")
    report.append("## 项目概述")
    report.append("")
    report.append("本项目旨在使用神经网络模型预测材料的太阳反射率和红外发射率，")
    report.append("基于材料的添加条件（高分子基体、无机添加剂、其余材料和加工方法）。")
    report.append("项目包括数据处理、模型训练、遗传算法优化和SHAP分析四个主要部分。")
    report.append("")
    
    # 添加数据处理部分
    report.append("## 数据处理")
    report.append("")
    report.append("数据处理部分包括以下步骤：")
    report.append("")
    report.append("1. 加载原始数据集")
    report.append("2. 清洗数据（处理缺失值、转换数据类型等）")
    report.append("3. 对分类特征进行独热编码")
    report.append("4. 将数据集分为训练集和测试集")
    report.append("5. 可视化数据分布和特征关系")
    report.append("")
    report.append("数据处理结果保存在 `output/` 目录中，可视化结果保存在 `figures/` 目录中。")
    report.append("")
    
    # 添加模型训练部分
    report.append("## 模型训练")
    report.append("")
    
    # 尝试加载最佳模型信息
    try:
        with open('output/best_model_info.json', 'r') as f:
            best_model_info = json.load(f)
        
        report.append(f"最佳模型: {best_model_info['name']}")
        report.append(f"太阳反射率 MSE: {best_model_info['solar_reflectance_mse']:.6f}")
        report.append(f"太阳反射率 R²: {best_model_info['solar_reflectance_r2']:.6f}")
        report.append(f"红外发射率 MSE: {best_model_info['infrared_emissivity_mse']:.6f}")
        report.append(f"红外发射率 R²: {best_model_info['infrared_emissivity_r2']:.6f}")
    except Exception as e:
        report.append("无法加载最佳模型信息。")
    
    report.append("")
    report.append("模型训练部分包括以下步骤：")
    report.append("")
    report.append("1. 实现了三种神经网络架构：MLP、残差网络和集成模型")
    report.append("2. 使用自定义损失函数（MSE + 相关性损失）训练模型")
    report.append("3. 实现了早停和学习率调度策略")
    report.append("4. 比较不同模型的性能并选择最佳模型")
    report.append("")
    report.append("模型训练结果保存在 `models/` 目录中，训练过程可视化结果保存在 `figures/` 目录中。")
    report.append("")
    
    # 添加遗传算法部分
    report.append("## 遗传算法优化")
    report.append("")
    report.append("遗传算法优化部分使用NSGA-III多目标遗传算法寻找最佳参数组合，以优化太阳反射率和红外发射率。")
    report.append("")
    report.append("优化过程包括以下步骤：")
    report.append("")
    report.append("1. 定义优化问题（变量、目标、约束等）")
    report.append("2. 运行NSGA-III算法寻找非支配解")
    report.append("3. 分析最佳解和参数分布")
    report.append("4. 可视化优化过程和结果")
    report.append("")
    report.append("遗传算法优化结果保存在 `ga_results/` 目录中。")
    report.append("")
    
    # 添加SHAP分析部分
    report.append("## SHAP分析")
    report.append("")
    report.append("SHAP分析部分使用SHAP（SHapley Additive exPlanations）方法解释模型预测结果，")
    report.append("帮助理解不同特征对预测太阳反射率和红外发射率的影响。")
    report.append("")
    report.append("分析过程包括以下步骤：")
    report.append("")
    report.append("1. 计算测试集样本的SHAP值")
    report.append("2. 可视化特征重要性和依赖关系")
    report.append("3. 分析特征交互")
    report.append("4. 分析特定样本（最佳预测和最差预测）")
    report.append("")
    report.append("SHAP分析结果保存在 `shap_results/` 目录中。")
    report.append("")
    
    # 添加结论部分
    report.append("## 结论与建议")
    report.append("")
    report.append("根据模型训练、遗传算法优化和SHAP分析的结果，我们可以得出以下结论：")
    report.append("")
    report.append("1. 神经网络模型能够有效预测材料的太阳反射率和红外发射率")
    report.append("2. 遗传算法优化能够找到具有高太阳反射率和高红外发射率的材料组合")
    report.append("3. SHAP分析揭示了影响材料性能的关键因素")
    report.append("")
    report.append("基于以上结论，我们建议：")
    report.append("")
    report.append("1. 在材料设计中优先考虑SHAP分析中识别的关键特征")
    report.append("2. 使用遗传算法优化的结果指导新材料的开发")
    report.append("3. 进一步收集数据以改进模型性能")
    report.append("")
    
    # 添加系统架构图部分
    report.append("## 系统架构图")
    report.append("")
    report.append("```")
    report.append("                      +----------------+")
    report.append("                      |   原始数据集   |")
    report.append("                      +-------+--------+")
    report.append("                              |")
    report.append("                              v")
    report.append("+------------------+  +-------+--------+")
    report.append("|  数据可视化与分析  |<-|   数据处理    |")
    report.append("+------------------+  +-------+--------+")
    report.append("                              |")
    report.append("                              v")
    report.append("                      +-------+--------+")
    report.append("                      |   模型训练    |")
    report.append("                      +-------+--------+")
    report.append("                              |")
    report.append("                      +-------v--------+")
    report.append("                      |   模型评估    |")
    report.append("                      +-------+--------+")
    report.append("                             / \\")
    report.append("                            /   \\")
    report.append("               +-----------+     +------------+")
    report.append("               |                              |")
    report.append("               v                              v")
    report.append("      +--------+-------+            +---------+------+")
    report.append("      | 遗传算法优化   |            |   SHAP分析     |")
    report.append("      +--------+-------+            +---------+------+")
    report.append("               |                              |")
    report.append("               v                              v")
    report.append("      +--------+-------+            +---------+------+")
    report.append("      | 最佳参数组合   |            | 特征重要性分析  |")
    report.append("      +----------------+            +----------------+")
    report.append("```")
    report.append("")
    
    # 添加优化建议部分
    report.append("## 优化建议")
    report.append("")
    report.append("为了进一步改进系统性能，我们提出以下优化建议：")
    report.append("")
    report.append("### 数据处理优化")
    report.append("")
    report.append("1. 收集更多样本数据，特别是对于稀有的材料组合")
    report.append("2. 考虑使用更高级的特征工程技术，如主成分分析（PCA）或自编码器")
    report.append("3. 实现更复杂的数据增强策略，如SMOTE或ADASYN")
    report.append("")
    report.append("### 模型优化")
    report.append("")
    report.append("1. 尝试更复杂的神经网络架构，如Transformer或图神经网络")
    report.append("2. 实现贝叶斯优化来自动调整超参数")
    report.append("3. 考虑使用迁移学习或预训练模型")
    report.append("")
    report.append("### 遗传算法优化")
    report.append("")
    report.append("1. 实现自适应变异和交叉策略")
    report.append("2. 考虑使用更高级的多目标优化算法，如MOEA/D或SPEA2")
    report.append("3. 增加约束条件以确保解的可行性")
    report.append("")
    report.append("### SHAP分析优化")
    report.append("")
    report.append("1. 实现更高效的SHAP值计算方法，如Deep SHAP或Tree SHAP")
    report.append("2. 考虑使用其他可解释性方法，如LIME或Integrated Gradients")
    report.append("3. 开发交互式可视化工具以便更好地探索SHAP结果")
    report.append("")
    
    # 将报告内容写入文件
    with open('summary_report.md', 'w', encoding='utf-8') as f:
        f.write('\n'.join(report))
    
    print("总结报告已生成: summary_report.md")


def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='材料性能预测系统')
    parser.add_argument('--skip-data-processing', action='store_true', help='跳过数据处理步骤')
    parser.add_argument('--skip-model-training', action='store_true', help='跳过模型训练步骤')
    parser.add_argument('--skip-genetic-algorithm', action='store_true', help='跳过遗传算法优化步骤')
    parser.add_argument('--skip-shap-analysis', action='store_true', help='跳过SHAP分析步骤')
    parser.add_argument('--only-report', action='store_true', help='只生成总结报告')
    args = parser.parse_args()
    
    # 记录开始时间
    start_time = time.time()
    
    # 创建必要的目录
    create_directories()
    
    # 如果只生成报告，则跳过所有其他步骤
    if args.only_report:
        generate_summary_report()
        return
    
    # 运行数据处理
    if not args.skip_data_processing:
        run_data_processing()
    else:
        print("跳过数据处理步骤")
    
    # 运行模型训练
    if not args.skip_model_training:
        run_model_training()
    else:
        print("跳过模型训练步骤")
    
    # 运行遗传算法
    if not args.skip_genetic_algorithm:
        run_genetic_algorithm()
    else:
        print("跳过遗传算法优化步骤")
    
    # 运行SHAP分析
    if not args.skip_shap_analysis:
        run_shap_analysis()
    else:
        print("跳过SHAP分析步骤")
    
    # 生成总结报告
    generate_summary_report()
    
    # 计算总运行时间
    end_time = time.time()
    run_time = end_time - start_time
    
    print("\n" + "=" * 50)
    print(f"系统运行完成！总耗时: {run_time:.2f} 秒")
    print("=" * 50)
    print("\n结果保存在以下目录:")
    print("- 数据处理结果: output/")
    print("- 模型文件: models/")
    print("- 可视化结果: figures/")
    print("- 遗传算法结果: ga_results/")
    print("- SHAP分析结果: shap_results/")
    print("- 总结报告: summary_report.md")


if __name__ == "__main__":
    main()